# -*- encoding: utf-8 -*-
"""
@File   : compare_data.py
@Author : l00464502
@Date   : 2022-7-16
"""
import os
import copy
import yaml

from log import WriteLog


class YamlHandle:
    """
    读写yaml文件
    """

    def __init__(self, file):
        self.file = file

    def read_yaml(self):
        """
        function: 读yaml文件数据
        :return: 字典或列表对象
        """
        with open(self.file, encoding='utf-8') as f:
            return yaml.load(f.read(), Loader=yaml.FullLoader)

    def write_yaml(self, data, encoding='utf-8'):
        """
        function: 写/修改yaml文件数据
        :return: 将python对象生成为yaml文档
        """
        with open(self.file, encoding=encoding, mode='w') as f:
            return yaml.dump(data, stream=f, allow_unicode=True)


class CompareData:
    """
    检查升级前后库内文件的一致性
    """

    def __init__(self, **kwargs):
        self.init_user = kwargs.get("init_user")
        self.db_user = kwargs.get("db_user")
        self.password = kwargs.get("password")
        self.database = kwargs.get("database")
        self.port = kwargs.get("port")
        self.env_path = kwargs.get("env_path")
        self.db_version = kwargs.get("db_version")
        self.table_list = []
        self.checksum_dict = {}
        self.check_ok = False
        self.log = WriteLog()
        self.check_yaml_file()

    @staticmethod
    def execute_command(command):
        """
        fucntion: 通过os.popen执行命令
        :param command: 需要执行的命令
        :return: 返回命令执行后的回显，列表形式
        """
        with os.popen(command, 'r') as p:
            f = p.readlines()
        return f

    def check_yaml_file(self):
        checksum_file_path = os.path.join(os.getcwd(), "log", "compare",
                                          self.db_version, "checksum.yaml")
        self.yaml_object = YamlHandle(checksum_file_path)
        if not os.path.exists(checksum_file_path):
            os.makedirs(os.path.dirname(checksum_file_path), exist_ok=True)
            initialize_data = {
                "checksum": {
                    "before_upgrade": "",
                    "after_upgrade": ""
                },
                "check_ok": False
            }
            self.yaml_object.write_yaml(initialize_data)

        self.data = self.yaml_object.read_yaml()
        return checksum_file_path

    def get_checksum(self):
        """
        function: 获取数据库各表校验和
        :return:
        """
        self.checksum_dict.clear()
        for table in self.table_list:
            gsql_command = f"su - {self.init_user} <<EOF\n" \
                           f"source {self.env_path};" \
                           f"gsql -d {self.database} -U {self.db_user} " \
                           f"-W {self.password} -p {self.port} " \
                           f"-c 'select checksum(\"{table}\"::text) " \
                           f"from \"{table}\";'" \
                           f" | awk 'NR==3 {{print \$1}}';\n" \
                           f"EOF"
            self.log.print_log(gsql_command)
            table_checksum = self.execute_command(gsql_command)[0].strip()
            self.checksum_dict.update({table: table_checksum})

    def write_checksum_to_file(self, key):
        """
        function: 将校验和字典写入文件
        :param key: 值为before_upgrade或after_upgrade
        :return:
        """
        try:
            self.data["checksum"][key] = copy.deepcopy(self.checksum_dict)
            self.data["check_ok"] = self.check_ok
            self.yaml_object.write_yaml(self.data)
        except Exception as e:
            self.log.print_error("出错啦！ {}".format(e))

    def get_tables(self):
        """
        function: 获取数据库中的表
        :return:
        """
        self.table_list.clear()
        gsql_command = f"su - {self.init_user} <<EOF\n" \
                       f"source {self.env_path};" \
                       f"gsql -d {self.database} -U {self.db_user} " \
                       f"-W {self.password} -p {self.port} " \
                       f"-c \"\dt+\" | awk 'NR>=3 && NF>2 {{print \$3}}';\n" \
                       f"EOF"
        self.log.print_log(gsql_command)
        self.table_list = \
            [x.strip() for x in self.execute_command(gsql_command)[:-1]]

    def compare(self):
        """
        function: 对比升级前后表的一致性
        :return:
        """
        data = self.yaml_object.read_yaml()
        before_upgrade = data["checksum"]["before_upgrade"]
        after_upgrade = data["checksum"]["after_upgrade"]

        diff_list = []
        diff_path = os.path.join(os.getcwd(), "log", "compare",
                                 self.db_version, "data.diff")

        if len(before_upgrade) != len(after_upgrade):
            self.log.print_error("注意：升级前后表数量不一致！")

        context = ""
        if before_upgrade == after_upgrade:
            self.check_ok = True
            self.log.write_file(diff_path, context)
            with open(diff_path, 'r+') as f:
                f.truncate(0)
            return

        self.log.print_error("注意：升级前后表不一致！")
        self.check_ok = False
        # 将不一致信息写入到diff文件
        before_upgrade_tables = before_upgrade.keys()
        after_upgrade_tables = after_upgrade.keys()

        diff_list.append("1.升级后多余表或者缺少表\n")
        tables_diff1 = set(before_upgrade_tables) - \
                       set(after_upgrade_tables)
        tables_diff2 = set(after_upgrade_tables) - \
                       set(before_upgrade_tables)
        for table in list(tables_diff1):
            diff_list.append({"before_upgrade": table, "after_grade": ""})
        for table in list(tables_diff2):
            diff_list.append({"before_upgrade": "", "after_grade": table})

        diff_list.append("2.升级后表内容不一致\n")
        both_have_tables = set(before_upgrade_tables) \
                           & set(after_upgrade_tables)
        for table in list(both_have_tables):
            if before_upgrade[table] != after_upgrade[table]:
                diff_list.append(f"\n{table}")
                diff_list.append({"before_upgrade": before_upgrade[table],
                                  "after_grade": after_upgrade[table]})
        context = "\n".join([str(x) for x in diff_list])
        self.log.write_file(diff_path, context)


if __name__ == "__main__":
    test = CompareData(init_user="init_user",
                       db_user="db_user",
                       password="password",
                       port=5432,
                       env_path="env_path",
                       database="database",
                       db_version="2.0.0->3.0.0")
    # 升级前取一次
    test.get_tables()
    test.get_checksum()
    test.write_checksum_to_file('before_upgrade')
    # 升级后取一次
    test.get_tables()
    test.get_checksum()
    test.write_checksum_to_file('after_upgrade')

    test.compare()
